Skip to content

[mlir][ArmSME] Prefix SME streaming attributes with "llvm." for propagation#190864

Merged
joker-eph merged 1 commit intollvm:mainfrom
joker-eph:fix_discardable_attr_sme
Apr 8, 2026
Merged

[mlir][ArmSME] Prefix SME streaming attributes with "llvm." for propagation#190864
joker-eph merged 1 commit intollvm:mainfrom
joker-eph:fix_discardable_attr_sme

Conversation

@joker-eph
Copy link
Copy Markdown
Collaborator

Function-level attributes enabling SME in the backend (e.g. arm_streaming, arm_locally_streaming) were dropped when converting func.func to llvm.func after ed37bdc, because attribute propagation during FuncToLLVM conversion only preserves attributes prefixed with "llvm.". Prefix these discardable attributes accordingly.

@llvmbot
Copy link
Copy Markdown
Member

llvmbot commented Apr 7, 2026

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

Changes

Function-level attributes enabling SME in the backend (e.g. arm_streaming, arm_locally_streaming) were dropped when converting func.func to llvm.func after ed37bdc, because attribute propagation during FuncToLLVM conversion only preserves attributes prefixed with "llvm.". Prefix these discardable attributes accordingly.


Full diff: https://github.com/llvm/llvm-project/pull/190864.diff

2 Files Affected:

  • (modified) mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp (+3-1)
  • (modified) mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir (+13-13)
diff --git a/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp b/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
index eafdc1de5ef34..8d14bff3b5875 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
@@ -128,7 +128,9 @@ struct EnableArmStreamingPass
 
     auto unitAttr = UnitAttr::get(&getContext());
 
-    function->setAttr(stringifyArmStreamingMode(streamingMode), unitAttr);
+    function->setDiscardableAttr(
+        (Twine("llvm.") + stringifyArmStreamingMode(streamingMode)).str(),
+        unitAttr);
 
     // The pass currently only supports enabling ZA when in streaming-mode, but
     // ZA can be accessed by the SME LDR, STR and ZERO instructions when not in
diff --git a/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir b/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
index 00b38b86f0d6e..087cc98edbdd6 100644
--- a/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
+++ b/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
@@ -6,42 +6,42 @@
 // RUN: mlir-opt %s -enable-arm-streaming=if-scalable-and-supported -verify-diagnostics | FileCheck %s -check-prefix=IF-SCALABLE
 
 // CHECK-LABEL: @arm_streaming
-// CHECK-SAME: attributes {arm_streaming}
+// CHECK-SAME: attributes {llvm.arm_streaming}
 // CHECK-LOCALLY-LABEL: @arm_streaming
-// CHECK-LOCALLY-SAME: attributes {arm_locally_streaming}
+// CHECK-LOCALLY-SAME: attributes {llvm.arm_locally_streaming}
 // CHECK-COMPATIBLE-LABEL: @arm_streaming
-// CHECK-COMPATIBLE-SAME: attributes {arm_streaming_compatible}
+// CHECK-COMPATIBLE-SAME: attributes {llvm.arm_streaming_compatible}
 // CHECK-ENABLE-ZA-LABEL: @arm_streaming
-// CHECK-ENABLE-ZA-SAME: attributes {arm_new_za, arm_streaming}
+// CHECK-ENABLE-ZA-SAME: attributes {llvm.arm_new_za, llvm.arm_streaming}
 func.func @arm_streaming() { return }
 
 // CHECK-LABEL: @not_arm_streaming
-// CHECK-SAME: attributes {enable_arm_streaming_ignore}
+// CHECK-SAME: attributes {llvm.enable_arm_streaming_ignore}
 // CHECK-LOCALLY-LABEL: @not_arm_streaming
-// CHECK-LOCALLY-SAME: attributes {enable_arm_streaming_ignore}
+// CHECK-LOCALLY-SAME: attributes {llvm.enable_arm_streaming_ignore}
 // CHECK-COMPATIBLE-LABEL: @not_arm_streaming
-// CHECK-COMPATIBLE-SAME: attributes {enable_arm_streaming_ignore}
+// CHECK-COMPATIBLE-SAME: attributes {llvm.enable_arm_streaming_ignore}
 // CHECK-ENABLE-ZA-LABEL: @not_arm_streaming
-// CHECK-ENABLE-ZA-SAME: attributes {enable_arm_streaming_ignore}
-func.func @not_arm_streaming() attributes {enable_arm_streaming_ignore} { return }
+// CHECK-ENABLE-ZA-SAME: attributes {llvm.enable_arm_streaming_ignore}
+func.func @not_arm_streaming() attributes {llvm.enable_arm_streaming_ignore} { return }
 
 // CHECK-LABEL: @requires_arm_streaming
-// CHECK-SAME: attributes {arm_streaming}
+// CHECK-SAME: attributes {llvm.arm_streaming}
 // IF-REQUIRED: @requires_arm_streaming
-// IF-REQUIRED-SAME: attributes {arm_streaming}
+// IF-REQUIRED-SAME: attributes {llvm.arm_streaming}
 func.func @requires_arm_streaming() {
   %tile = arm_sme.get_tile : vector<[4]x[4]xi32>
   return
 }
 
 // CHECK-LABEL: @does_not_require_arm_streaming
-// CHECK-SAME: attributes {arm_streaming}
+// CHECK-SAME: attributes {llvm.arm_streaming}
 // IF-REQUIRED: @does_not_require_arm_streaming
 // IF-REQUIRED-NOT: arm_streaming
 func.func @does_not_require_arm_streaming() { return }
 
 // IF-SCALABLE-LABEL: @contains_scalable_vectors
-// IF-SCALABLE-SAME: attributes {arm_streaming}
+// IF-SCALABLE-SAME: attributes {llvm.arm_streaming}
 func.func @contains_scalable_vectors(%vec: vector<[4]xf32>) -> vector<[4]xf32> {
   %0 = arith.addf %vec, %vec : vector<[4]xf32>
   return %0 : vector<[4]xf32>

@llvmbot
Copy link
Copy Markdown
Member

llvmbot commented Apr 7, 2026

@llvm/pr-subscribers-mlir-sme

Author: Mehdi Amini (joker-eph)

Changes

Function-level attributes enabling SME in the backend (e.g. arm_streaming, arm_locally_streaming) were dropped when converting func.func to llvm.func after ed37bdc, because attribute propagation during FuncToLLVM conversion only preserves attributes prefixed with "llvm.". Prefix these discardable attributes accordingly.


Full diff: https://github.com/llvm/llvm-project/pull/190864.diff

2 Files Affected:

  • (modified) mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp (+3-1)
  • (modified) mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir (+13-13)
diff --git a/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp b/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
index eafdc1de5ef34..8d14bff3b5875 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
@@ -128,7 +128,9 @@ struct EnableArmStreamingPass
 
     auto unitAttr = UnitAttr::get(&getContext());
 
-    function->setAttr(stringifyArmStreamingMode(streamingMode), unitAttr);
+    function->setDiscardableAttr(
+        (Twine("llvm.") + stringifyArmStreamingMode(streamingMode)).str(),
+        unitAttr);
 
     // The pass currently only supports enabling ZA when in streaming-mode, but
     // ZA can be accessed by the SME LDR, STR and ZERO instructions when not in
diff --git a/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir b/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
index 00b38b86f0d6e..087cc98edbdd6 100644
--- a/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
+++ b/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
@@ -6,42 +6,42 @@
 // RUN: mlir-opt %s -enable-arm-streaming=if-scalable-and-supported -verify-diagnostics | FileCheck %s -check-prefix=IF-SCALABLE
 
 // CHECK-LABEL: @arm_streaming
-// CHECK-SAME: attributes {arm_streaming}
+// CHECK-SAME: attributes {llvm.arm_streaming}
 // CHECK-LOCALLY-LABEL: @arm_streaming
-// CHECK-LOCALLY-SAME: attributes {arm_locally_streaming}
+// CHECK-LOCALLY-SAME: attributes {llvm.arm_locally_streaming}
 // CHECK-COMPATIBLE-LABEL: @arm_streaming
-// CHECK-COMPATIBLE-SAME: attributes {arm_streaming_compatible}
+// CHECK-COMPATIBLE-SAME: attributes {llvm.arm_streaming_compatible}
 // CHECK-ENABLE-ZA-LABEL: @arm_streaming
-// CHECK-ENABLE-ZA-SAME: attributes {arm_new_za, arm_streaming}
+// CHECK-ENABLE-ZA-SAME: attributes {llvm.arm_new_za, llvm.arm_streaming}
 func.func @arm_streaming() { return }
 
 // CHECK-LABEL: @not_arm_streaming
-// CHECK-SAME: attributes {enable_arm_streaming_ignore}
+// CHECK-SAME: attributes {llvm.enable_arm_streaming_ignore}
 // CHECK-LOCALLY-LABEL: @not_arm_streaming
-// CHECK-LOCALLY-SAME: attributes {enable_arm_streaming_ignore}
+// CHECK-LOCALLY-SAME: attributes {llvm.enable_arm_streaming_ignore}
 // CHECK-COMPATIBLE-LABEL: @not_arm_streaming
-// CHECK-COMPATIBLE-SAME: attributes {enable_arm_streaming_ignore}
+// CHECK-COMPATIBLE-SAME: attributes {llvm.enable_arm_streaming_ignore}
 // CHECK-ENABLE-ZA-LABEL: @not_arm_streaming
-// CHECK-ENABLE-ZA-SAME: attributes {enable_arm_streaming_ignore}
-func.func @not_arm_streaming() attributes {enable_arm_streaming_ignore} { return }
+// CHECK-ENABLE-ZA-SAME: attributes {llvm.enable_arm_streaming_ignore}
+func.func @not_arm_streaming() attributes {llvm.enable_arm_streaming_ignore} { return }
 
 // CHECK-LABEL: @requires_arm_streaming
-// CHECK-SAME: attributes {arm_streaming}
+// CHECK-SAME: attributes {llvm.arm_streaming}
 // IF-REQUIRED: @requires_arm_streaming
-// IF-REQUIRED-SAME: attributes {arm_streaming}
+// IF-REQUIRED-SAME: attributes {llvm.arm_streaming}
 func.func @requires_arm_streaming() {
   %tile = arm_sme.get_tile : vector<[4]x[4]xi32>
   return
 }
 
 // CHECK-LABEL: @does_not_require_arm_streaming
-// CHECK-SAME: attributes {arm_streaming}
+// CHECK-SAME: attributes {llvm.arm_streaming}
 // IF-REQUIRED: @does_not_require_arm_streaming
 // IF-REQUIRED-NOT: arm_streaming
 func.func @does_not_require_arm_streaming() { return }
 
 // IF-SCALABLE-LABEL: @contains_scalable_vectors
-// IF-SCALABLE-SAME: attributes {arm_streaming}
+// IF-SCALABLE-SAME: attributes {llvm.arm_streaming}
 func.func @contains_scalable_vectors(%vec: vector<[4]xf32>) -> vector<[4]xf32> {
   %0 = arith.addf %vec, %vec : vector<[4]xf32>
   return %0 : vector<[4]xf32>

…gation

Function-level attributes enabling SME in the backend (e.g.
arm_streaming, arm_locally_streaming) were dropped when converting
func.func to llvm.func after ed37bdc, because attribute propagation
during FuncToLLVM conversion only preserves attributes prefixed with
"llvm.". Prefix these discardable attributes accordingly.
@joker-eph joker-eph force-pushed the fix_discardable_attr_sme branch from 4a5ffd4 to c170ed8 Compare April 7, 2026 22:36
@hockyy
Copy link
Copy Markdown
Contributor

hockyy commented Apr 8, 2026

lgtm

Copy link
Copy Markdown
Contributor

@banach-space banach-space left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks!

Btw, I am working on introducing a dedicated buildbot for this sort of stuff. As usual, these things take time.

@luporl
Copy link
Copy Markdown
Contributor

luporl commented Apr 8, 2026

Thanks for this fix. On my AArch64 machine it fixed many tests, but 2 are still failing:

MLIR :: Integration/Dialect/Linalg/CPU/ArmSME/multi-tile-matmul-mixed-types.mlir

# | /home/leandro.lupori/git/ci/llvm-project/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/multi-tile-matmul-mixed-types.mlir:57:12: error: CHECK: expected string not found in input
# |  // CHECK: [32955, 33514, 34073, 34632, 35191, 35750, 36309]
# |            ^
# | <stdin>:1:98: note: scanning from here
# | Unranked Memref base@ = 0xaaaab0a63780 rank = 2 offset = 0 sizes = [7, 7] strides = [7, 1] data =
# |                                                                                                  ^
# | <stdin>:2:2: note: possible intended match here
# | [[32955, 34086, 34073, 35230, 35191, 36374, 36309],
# |  ^
...

MLIR :: Integration/Dialect/Vector/CPU/ArmSME/outerproduct-i8i8i32.mlir

# | /home/leandro.lupori/git/ci/llvm-project/mlir/test/Integration/Dialect/Vector/CPU/ArmSME/outerproduct-i8i8i32.mlir:65:12: error: CHECK: expected string not found in input
# |  // CHECK: ( 110, 134, 158, 182 )
# |            ^
# | <stdin>:1:1: note: scanning from here
# | ( 110, 478, 158, 654 )
# | ^
...

@joker-eph
Copy link
Copy Markdown
Collaborator Author

Let me merge this and follow-up on these then.

@joker-eph joker-eph merged commit 7765fd8 into llvm:main Apr 8, 2026
10 checks passed
@joker-eph joker-eph deleted the fix_discardable_attr_sme branch April 8, 2026 20:13
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants